#!/usr/bin/python
# -*- coding: utf-8 -*-
from optparse import OptionParser
import pywintypes
import sys
import trace
import traceback
import win32api
import win32process

#-----------------------------------------------------------------------------
# Name:        
# Purpose:     
#
# Author:      Administrator
#
# Created:     2006-12-5 下午05:32:15
# RCS-ID:      Id:  
# Copyright:   (c) 2006
# Licence:     Free To Use
# 
#-----------------------------------------------------------------------------

#import win32ui
#import win32con

def runner(p_aff, p_n):
    try:
        p_handle = win32process.CreateProcess(None, p_n, None , None , 0 , win32process.NORMAL_PRIORITY_CLASS , None , None , win32process.STARTUPINFO())
        print >> sys.stdout, 'Running Process Success'
    except:
        print >> sys.stderr, 'Runnning Process Error'
        print sys.exc_info()[1][2]
        sys.exit()
        
    try:
        win32process.SetProcessAffinityMask(p_handle[0], p_aff)
        __CURRENT_AFF = win32process.GetProcessAffinityMask(p_handle[0])[0]
        print >> sys.stdout, 'Set Process CPU Affinity Mask Sucess'
        print >> sys.stdout, "Current Process running as CPU Affinity Mask %d" % __CURRENT_AFF
    except:
        print >> sys.stderr, 'Set Process CPU Affinity Mask Error'
        print sys.exc_info()[1][2]
    
def main():
    __CPU_INFO = win32api.GetSystemInfo()
    __CPU_NUMBER = __CPU_INFO[5]
    __CPU_DEF_AFF_MASK = __CPU_INFO[4]
    
    print >> sys.stdout, """
Current Number Of CPU :%d 
Current Default CPU Affinity Mask : %d
    """ % (__CPU_NUMBER, __CPU_DEF_AFF_MASK)
    
    gl_usage = "usage: %prog [options1] arg1 [options2] arg2"
    my_opt = OptionParser(usage=gl_usage, version="%prog 0.1")
    my_opt.prog = "setaff.exe"
    my_opt.description = """Run Process And Set Process Affinity Mask. """
    my_opt.add_option('-c', '--cpu', dest='cpu', \
                      type="int", default=1, \
                      help='Running Process On that CPU ID. 1=CPU1 2=CPU2 3=BOTH. default running on CPU 1')
    my_opt.add_option('-p', '--process', dest='process', \
                      type="string", default=None, help='Running Process command line')

    (in_options, in_args) = my_opt.parse_args()
    
    if in_options.process == None:
        print >> sys.stderr, """Need a command line to running

usage: setaff.exe [options1] arg1 [options2] arg2

Run Process And Set Process Affinity Mask

options:
  --version             show program's version number and exit
  -h, --help            show this help message and exit
  -c CPU, --cpu=CPU     Running Process On that CPU ID. 1=CPU1 2=CPU2 3=BOTH.
                        default running on CPU 1
  -p PROCESS, --process=PROCESS
                        Running Process command line
                        """
        sys.exit()
    elif in_options.cpu == None and in_options.process != None:
        __CPU__ = 1
        print >> sys.stdout, 'Running Process on CPU 1'
        runner(__CPU__, in_options.process)
    elif in_options.cpu != None and in_options.process != None:
        print >> sys.stdout, 'Try to Running Process on CPU %s' % in_options.cpu
        runner(in_options.cpu, in_options.process)
    else:
        print >> sys.stderr, 'Some Thing Wrong'
        sys.exit()

if __name__ == '__main__':
    main()